import java.io.File;
import java.io.IOException;
import java.lang.management.ManagementFactory;
import java.lang.management.MemoryMXBean;
import java.lang.management.MemoryPoolMXBean;
import java.lang.management.MemoryUsage;

import javax.management.MBeanServerConnection;
import javax.management.ObjectName;
import javax.management.remote.JMXConnectorFactory;
import javax.management.remote.JMXServiceURL;

import com.sun.tools.attach.AgentInitializationException;
import com.sun.tools.attach.AgentLoadException;
import com.sun.tools.attach.VirtualMachine;
import com.sun.tools.attach.VirtualMachineDescriptor;

public class HeapInfo {
	private static final String LOCAL_CONNECTOR_ADDRESS_PROP =
		"com.sun.management.jmxremote.localConnectorAddress";

	public static void main(String... args) {
		int n = args.length;
		if (n == 0 || "-all".equals(args[0])) {
			for (VirtualMachineDescriptor desc : VirtualMachine.list()) {
				String pid = desc.id();
				System.out.println(pid + "\t" + desc.displayName());
				process(pid);
				System.out.println();
			}
		} else if ("-pid".equals(args[0])) {
			for (int i = 1; i < n; ++i) {
				String pid = args[i];
				System.out.println(pid);
				process(pid);
			}
		} else {
			System.out.println("Arguments:");
			System.out.println("\t-all            process all local virtual machines");
			System.out.println("\t-pid <p0>...    process local virtual machines by their pids");
		}
	}

	static void process(String pid) {
		try {
			VirtualMachine vm = VirtualMachine.attach(pid);
			try {
				String address = getConnectorAddress(vm);
				if (address != null) {
					MBeanServerConnection connection = JMXConnectorFactory.connect(new JMXServiceURL(address)).getMBeanServerConnection();
					MemoryMXBean memoryMXBean = ManagementFactory.newPlatformMXBeanProxy(connection, ManagementFactory.MEMORY_MXBEAN_NAME, MemoryMXBean.class);
					MemoryUsage hp = memoryMXBean.getHeapMemoryUsage();
					MemoryUsage nh = memoryMXBean.getNonHeapMemoryUsage();
					printMemoryUsage("Heap", hp);
					printMemoryUsage("NonHeap", nh);
					for (ObjectName name : connection.queryNames(new ObjectName(ManagementFactory.MEMORY_POOL_MXBEAN_DOMAIN_TYPE + ",name=*"), null)) {
						MemoryPoolMXBean poolMXBean = ManagementFactory.newPlatformMXBeanProxy(connection, name.getCanonicalName(), MemoryPoolMXBean.class);
						String pool = poolMXBean.getName();
						MemoryUsage us = poolMXBean.getUsage();
						printMemoryUsage(pool, us);
					}
				}
			} finally {
				vm.detach();
			}
		} catch (Exception exc) {
			System.out.println(exc);
		}
	}

	static void printMemoryUsage(String name, MemoryUsage usage) {
		final double M = 1024 * 1024, K = 1024;
		System.out.printf("%20s\t%11.3f\t%11.3f\t%11.3f\t%11.3f\n", name, usage.getInit() / M, usage.getUsed() / M, usage.getCommitted() / M, usage.getMax() / M);
		System.out.printf("%20s\t%11.3f\t%11.3f\t%11.3f\t%11.3f\n",   "", usage.getInit() / K, usage.getUsed() / K, usage.getCommitted() / K, usage.getMax() / K);
	}

	static String getConnectorAddress(VirtualMachine vm) throws IOException, AgentInitializationException, AgentLoadException {
		String address = vm.getAgentProperties().getProperty(LOCAL_CONNECTOR_ADDRESS_PROP);
		if (address == null) {
			String home = vm.getSystemProperties().getProperty("java.home");
			// Normally in ${java.home}/jre/lib/management-agent.jar but might
			// be in ${java.home}/lib in build environments.
			File agent = new File(home + File.separator + "jre" + File.separator + "lib" + File.separator + "management-agent.jar");
			if (!agent.exists()) {
				agent = new File(home + File.separator +  "lib" + File.separator + "management-agent.jar");
				if (!agent.exists()) {
					agent = null;
				}
			}
			if (agent != null) {
				vm.loadAgent(agent.getCanonicalPath(), "com.sun.management.jmxremote");
			}
			address = vm.getAgentProperties().getProperty(LOCAL_CONNECTOR_ADDRESS_PROP);
		}
		return address;
	}
}
